#!/bin/bash
#
# Flash the NVMe disk with the system.img.raw
#

DISKDEV=
DISKSIZE=
BOARD=
ROOTFSIMG="./bootloader/system.img.raw"

MNTPOINT_IMG=/mnt/sys_image
MNTPOINT_DISK=/mnt/nvme_disk

OPTSTR="d:f:g:h"

function usage()
{
	echo -e "Usage:"
	echo -e "  $0 <opts>"
	echo -e "    -d <NVMeDisk> ----- The NVMe Disk device node in the host, can be: sdb, sdc, ..."
	echo -e "    -f <RootfsImg> ---- The rootfs image file path, default is: ${ROOTFSIMG}"
	echo -e "    -g <Board> -------- Generate image from L4T with the given board configuration"
	echo -e "    -h ---------------- Show this usage"
	echo -e "Examples:"
	echo -e "  $0 -d <NVMeDisk> [-f <RootfsImg>]"
	echo -e "  $0 -d <NVMeDisk> [-f <RootfsImg>]"
	echo -e "  $0 -d <NVMeDisk> -g <Board>\n"
}

function get_options()
{
	ARGS=$(getopt -o ${OPTSTR} -n "$0" -- "$@")
	if [[ $? != 0 ]];
	then
		usage
		exit 1
	fi

	eval set -- ${ARGS}

	# -g, -f options conflict
	OPT_CONFLICT=$(echo "${ARGS}" | grep -Eo '\-g|\-f' | wc -l)
	if [[ ${OPT_CONFLICT} == "2" ]]; then
		echo "[ERROR] Options conflict: -g, -f"
		usage
		exit 1
	fi

	while true; do
		case "$1" in
		-d)
			DISKDEV=$2
			shift 2
			;;
		-g)
			BOARD=$2
			shift 2
			;;
		-f)
			ROOTFSIMG=$2
			shift 2
			;;
		-h)
			usage
			exit 0
			;;
		--)
			break
			;;
		*)
			echo "[ERROR] Unsupport options: $2"
			usage
			exit 1
			;;
		esac
	done

	if [[ -z ${DISKDEV} ]]; then
		echo -e "[ERROR] NVMeDEV must be specified.\n"
		usage
		exit 1
	else
		if [[ ! ${DISKDEV} =~ ^sd[b-z]$ ]]; then
			echo -e "[ERROR] NVMe device name error, should be \"sdb, sdc, ...\".\n"
			exit 1
		fi
		if [[ ! -b /dev/${DISKDEV} ]]; then
			echo -e "[ERROR] NVMe device node \"/dev/${DISKDEV}\" does not exist.\n"
			exit 1
		fi
	fi

	if [[ -n ${BOARD} ]]; then
		if [[ ! -f ${BOARD}.conf ]]; then
			echo -e "[ERROR] Board configuration file does not exist: ${BOARD}.conf\n"
			exit 1
		fi
	else
		if [[ ! -f ${ROOTFSIMG} ]]; then
			echo -e "[ERROR] Rootfs image does not exist: ${ROOTFSIMG}"
			exit 1
		fi
	fi
}

function format_disk()
{
	# Unmount first if already mounted
	lsblk -nro NAME,MOUNTPOINT /dev/${DISKDEV} | while read line
	do
		partition=$(echo ${line} | awk {'print $1'})
		mountpoint=$(echo ${line} | awk {'print $2'})
		if [[ -n ${mountpoint} ]] && [[ ${mountpoint} =~ ^/ ]]; then
			echo -e "\n[INFO] Umount disk partition:"
			echo -e ">>> ${partition} ${mountpoint}"
			sudo umount ${mountpoint}
		fi
	done

	echo -e "[INFO] Create GPT\n"
	sudo parted -s /dev/${DISKDEV} mklabel gpt
	RET=$?
	if [[ ${RET} -ne 0 ]]; then
		echo -e "[ERROR] Create GPT error.\n"
		exit 1
	fi

	echo -e "[INFO] Add APP partition\n"
	sudo parted -s /dev/${DISKDEV} mkpart APP 0% 100%
	RET=$?
	if [[ ${RET} -ne 0 ]]; then
		echo -e "[ERROR] Add APP partition error.\n"
		exit 1
	fi

	echo -e "[INFO] Format APP partition as an ext4 filesystem\n"
	sudo mkfs -t ext4 /dev/${DISKDEV}1
}

function generate_image_from_rootfs()
{
	if [[ ! -f flash.sh ]]; then
		echo -e "[ERROR] Please change to <L4T_TOP> directory.\n"
		exit 1
	fi

	sudo BOOTDEV=nvme0n1p1 ./flash.sh --no-flash ${BOARD} nvme0n1p1
	RET=$?
	if [[ ${RET} -ne 0 ]]; then
		echo -e "[ERROR] Generating rootfs image error.\n"
		exit 1
	fi
}

function copy_from_image_to_nvme()
{
	if [[ -d ${MNTPOINT_IMG} ]]; then
		local mountlist=$(findmnt -n ${MNTPOINT_IMG} | wc -l)
		if [[ ${mountlist} -gt 0 ]]; then
			sudo umount ${MNTPOINT_IMG}
		fi
	else
		sudo mkdir ${MNTPOINT_IMG}
	fi

	if [[ -d ${MNTPOINT_DISK} ]]; then
		mountlist=$(findmnt -n ${MNTPOINT_DISK} | wc -l)
		if [[ ${mountlist} -gt 0 ]]; then
			sudo umount ${MNTPOINT_DISK}
		fi
	else
		sudo mkdir ${MNTPOINT_DISK}
	fi

	# Copy rootfs files into disk
	sudo mount ${ROOTFSIMG} ${MNTPOINT_IMG}
	RET=$?
	if [[ ${RET} -ne 0 ]]; then
		echo -e "[ERROR] Mount ${ROOTFSIMG} error.\n"
		exit 1
	fi
	sudo mount /dev/${DISKDEV}1 ${MNTPOINT_DISK}
	RET=$?
	if [[ ${RET} -ne 0 ]]; then
		echo -e "[ERROR] Mount /dev/${DISKDEV}1 error.\n"
		exit 1
	fi
	sudo rsync -axHAWX --numeric-ids --info=progress2 --exclude=/proc ${MNTPOINT_IMG}/ ${MNTPOINT_DISK}

	# Unmount the flash drive and rootfs image
	sudo umount ${MNTPOINT_IMG} ${MNTPOINT_DISK}
}

####################################################################

if [[ $(id -u) -ne 0 ]]; then
	echo -e "[ERROR] Please run as root!\n"
	exit 1
fi

get_options $@

DISKSIZE=$(lsblk -o NAME,SIZE /dev/${DISKDEV} | grep ^${DISKDEV} | awk {'print $2'})
echo "DISKDEV=${DISKDEV}"
echo "DISKSIZE=${DISKSIZE}"
echo "BOARD=${BOARD}"
echo "ROOTFSIMG=${ROOTFSIMG}"

echo -e "\n[INFO] Format NVMe disk device: ${DISKDEV}"
echo -e "[WARN] The disk data will lost!"
read -p "Continue (Y/n)?" key
if [[ ${key} == "N" || ${key} == "n" ]]; then
	echo -e "[INFO] Exit."
	exit 0
fi
format_disk

if [[ -n ${BOARD} ]]; then
	echo -e "\n[INFO] Generate the rootfs image from L4T"
	echo -e ">>> ${ROOTFSIMG}"
	generate_image_from_rootfs
fi

echo -e "\n[INFO] Copy rootfs files from ROOTFSIMG to DISKDEV ..."
echo -e ">>> From: ${ROOTFSIMG}"
echo -e ">>> To  : /dev/${DISKDEV}1"
copy_from_image_to_nvme

echo -e "\n[INFO] Done!\n"
